
from ship_envs.core.obs_viewer import ObsViewer
from ship_envs.core.observation import get_obs_img_size


class World(object):
    def __init__(self):
        self.agent_ships = []
        self.other_ships = []
        self.landmarks = []
        self.viewer = ObsViewer(*get_obs_img_size())
        self.obs_type = 0 # 0:图片，1:位置， 2:图片+位置

    def get_entitys(self):
        return self.agent_ships + self.other_ships + self.landmarks

    def get_agent_ships(self):
        return self.agent_ships

    def get_other_ships(self):
        return self.other_ships

    def get_landmarks(self):
        return self.landmarks

    def get_agent_ship(self, index=0):
        if self.agent_ships:
            return self.agent_ships[index]
        return None

    def get_other_ship(self, index=0):
        if self.other_ships:
            return self.other_ships[index]
        return None

    def step(self, action):
        for i, ship in enumerate(self.agent_ships):
            ship.step(action[i])
        for ship in self.other_ships:
            ship.step(0)
        self._add_models2viewer()

    def has_collision(self):
        """
        碰撞检测，检测船与船、船与障碍物之间是否发生碰撞
        :return:
        """
        # 检测船与船
        ships = self.agent_ships
        for ship in ships:
            for other in ships:
                if ship == other:
                    continue
                if ship.collision_with_other(other):
                    return True
        # 暂不考虑目标船的碰撞
        for agent_ship in self.agent_ships:
            for other in self.other_ships:
                if agent_ship.collision_with_other(other):
                    return True
        # 检测船与障碍物
        for agent_ship in self.agent_ships:
            for landmark in self.landmarks:
                if agent_ship.collision_with_other(landmark):
                    return True
        return False

    def _add_models2viewer(self):
        for landmark in self.landmarks:
            self.viewer.add_entity(landmark)
        for ship in self.other_ships:
            self.viewer.add_entity(ship)
        # 最后添加智能体船及目标，确保在最上层图层
        for ship in self.agent_ships:
            self.viewer.add_entity(ship)


    def add_ship(self, ship):
        if ship:
            if ship.is_agent():
                self.agent_ships.append(ship)
            else:
                self.other_ships.append(ship)

    def add_landmark(self, landmark):
        self.landmarks.append(landmark)

    def refresh(self):
        self.viewer.clear_models()
        self._add_models2viewer()

    def reset(self):
        self.step_count = 0

    def get_obs(self):
        if self.viewer:
            return self.viewer.get_array()
        else:
            return None


